# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import os
import datetime
import glob
import traceback

import pandas as pd

from absl import (flags, app)

FLAGS = flags.FLAGS

flags.DEFINE_string('log_dir_path', '', '')
flags.DEFINE_string('start_date', '', '')
flags.DEFINE_string('end_date', '', '')
flags.DEFINE_string('output_dir', 'plot/basis_smm', '')
flags.DEFINE_boolean('is_futures', False, '')


def find_all_balance_within(from_dt, to_dt, log_path, is_futures):
  details = []
  skip = True
  with open(log_path, 'r') as fi:
    in_log = False
    dt = None
    for line in fi:
      line = line.strip()
      if line.startswith('* ') or line.startswith('---  -'):
        if not in_log:
          try:
            fmt = '%Y-%m-%d %H:%M:%S.%f' if '.' in line else '%Y-%m-%d %H:%M:%S'
            dt = datetime.datetime.strptime(line.split(' ', 1)[1], fmt)
            skip = not (from_dt <= dt <= to_dt)
          except Exception:
            continue
          in_log = True
        else:
          if skip:
            continue
          items = [item for item in line[1:].strip().split(' ') if item]
          if is_futures:
            if not items[0].startswith('Total'):
              continue
            details.append({'datetime': dt, 'ticker': items[0][5:], 'balance': float(items[-1])})
          else:
            if items[0] in ['--', 'Currency', 'Total'] or len(items) <= 5:
              continue
            details.append({
                'datetime': dt,
                'ticker': items[0],
                'reserve': float(items[1]),
                'balance': float(items[2]),
                'QuotePrice': float(items[3])
            })
      else:
        in_log = False
  df = pd.DataFrame(details)
  return df


def calculate_and_output_pnl_breakdown(df, from_dt, to_dt, is_futures):
  from_dt = from_dt.strftime('%Y%m%d%H%M%S')
  to_dt = to_dt.strftime('%Y%m%d%H%M%S')
  for ticker, tdf in df.groupby('ticker'):
    if is_futures:
      tdf['%s_pl' % ticker] = tdf['balance'] - tdf.iloc[0]['balance']
    else:
      tdf['diff'] = tdf['balance'] - tdf['reserve']
      tdf['apl'] = tdf['diff'] * (tdf['QuotePrice'].shift(-1) - tdf['QuotePrice'])
      tdf['%s_pl' % ticker] = tdf['apl'].cumsum()
    strategy_name = FLAGS.log_dir_path
    if strategy_name.startswith('log_'):
      strategy_name = strategy_name[4:]
    tdf.set_index('datetime')[['%s_pl' % ticker]].plot(title=strategy_name)
    plt.savefig(
        os.path.join(FLAGS.output_dir,
                     '%s_pl_breakdown_%s_%s_%s' % (strategy_name, ticker, from_dt, to_dt)))


def main(argv):
  if not FLAGS.log_dir_path:
    print('please specify --log_dir_path')
    return 1
  logs = glob.glob('%s/*' % FLAGS.log_dir_path)
  assert len(logs) == 1, logs

  if FLAGS.end_date:
    end_date = datetime.datetime.strptime(FLAGS.end_date, '%Y%m%d')
  else:
    end_date = datetime.datetime.strptime(datetime.datetime.now().strftime('%Y%m%d'), '%Y%m%d')
  if FLAGS.start_date:
    start_date = datetime.datetime.strptime(FLAGS.start_date, '%Y%m%d')
  else:
    start_date = end_date
  is_futures = FLAGS.is_futures
  end_date += datetime.timedelta(days=1)
  try:
    df = find_all_balance_within(start_date, end_date, logs[0], is_futures)
    if not os.path.exists(FLAGS.output_dir):
      os.makedirs(FLAGS.output_dir)
    calculate_and_output_pnl_breakdown(df, start_date, end_date, is_futures)
  except Exception:
    print('%s falied' % FLAGS.log_dir_path)
    traceback.print_exc()


if __name__ == '__main__':
  app.run(main)
